{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z9ln8ylCKHsx"
   },
   "source": [
    "# L3-C - Linear Quantization II: Per Channel Quantization\n",
    "\n",
    "In this lesson, you will continue to learn about different granularities of performing linear quantization. You will cover `per channel` in this notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the next cell to import all of the functions you have used before in the previous lesson(s) of `Linear Quantization II` to follow along with the video.\n",
    "\n",
    "- To access the `helper.py` file, you can click `File --> Open...`, on the top left."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 97
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from helper import get_q_scale_symmetric, linear_q_with_scale_and_zero_point, linear_dequantization\n",
    "from helper import plot_quantization_errors, quantization_error"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JA1-rcLz4t4D"
   },
   "source": [
    "## Different Granularities for Quantization\n",
    "- For simplicity, you'll perform these using Symmetric mode."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oo4BCLpsDw3t"
   },
   "source": [
    "### Per Channel\n",
    "- Implement `Per Channel` Symmetric Quantization\n",
    "- `dim` parameter decides if it needs to be along the rows or columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 114
   },
   "outputs": [],
   "source": [
    "def linear_q_symmetric_per_channel(tensor,dim,dtype=torch.int8):\n",
    "\n",
    "\n",
    "\n",
    "    return quantized_tensor, scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 97
   },
   "outputs": [],
   "source": [
    "test_tensor=torch.tensor(\n",
    "    [[191.6, -13.5, 728.6],\n",
    "     [92.14, 295.5,  -184],\n",
    "     [0,     684.6, 245.5]]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- `dim = 0`, along the rows\n",
    "- `dim = 1`, along the columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 46
   },
   "outputs": [],
   "source": [
    "dim=0\n",
    "output_dim = test_tensor.shape[dim]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "output_dim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 129,
     "status": "ok",
     "timestamp": 1705361580592,
     "user": {
      "displayName": "Marc Sun",
      "userId": "00829270524676809963"
     },
     "user_tz": 300
    },
    "height": 29,
    "id": "tnEv2pJ8D56J"
   },
   "outputs": [],
   "source": [
    "scale = torch.zeros(output_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "scale"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Iterate through each row to calculate its `scale`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 122,
     "status": "ok",
     "timestamp": 1705361582194,
     "user": {
      "displayName": "Marc Sun",
      "userId": "00829270524676809963"
     },
     "user_tz": 300
    },
    "height": 80,
    "id": "kDMFi-enEcDC"
   },
   "outputs": [],
   "source": [
    "for index in range(output_dim):\n",
    "    sub_tensor = test_tensor.select(dim,index)\n",
    "    # print(sub_tensor)\n",
    "    scale[index] = get_q_scale_symmetric(sub_tensor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 425
    },
    "executionInfo": {
     "elapsed": 875,
     "status": "ok",
     "timestamp": 1705361583875,
     "user": {
      "displayName": "Marc Sun",
      "userId": "00829270524676809963"
     },
     "user_tz": 300
    },
    "height": 29,
    "id": "OHd3YWNBYfW3",
    "outputId": "bbd84be6-b78a-41c5-8b64-2f412d524834"
   },
   "outputs": [],
   "source": [
    "scale_shape = [1] * test_tensor.dim()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 425
    },
    "executionInfo": {
     "elapsed": 805,
     "status": "ok",
     "timestamp": 1705361586500,
     "user": {
      "displayName": "Marc Sun",
      "userId": "00829270524676809963"
     },
     "user_tz": 300
    },
    "height": 29,
    "id": "k4hkoyA7Ykve",
    "outputId": "3b4d185b-cdbe-4bbc-e748-d93d714d512b"
   },
   "outputs": [],
   "source": [
    "scale_shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "scale_shape[dim] = -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "scale_shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "scale = scale.view(scale_shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 80
   },
   "outputs": [],
   "source": [
    "# copy to be used later\n",
    "copy_scale = scale\n",
    "\n",
    "scale"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Understanding tensor by tensor division using `view` function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "m = torch.tensor([[1,2,3],[4,5,6],[7,8,9]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "s = torch.tensor([1,5,10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "s.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "s.view(1, 3).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 46
   },
   "outputs": [],
   "source": [
    "# alternate way\n",
    "s.view(1, -1).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "s.view(-1,1).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Along the row division"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "scale = torch.tensor([[1], [5], [10]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "scale.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "m / scale"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Along the column division"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "scale = torch.tensor([[1, 5, 10]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "scale.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "m / scale"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Coming back to quantizing the tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 80
   },
   "outputs": [],
   "source": [
    "# the scale you got earlier\n",
    "scale = copy_scale\n",
    "\n",
    "scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "scale.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 46
   },
   "outputs": [],
   "source": [
    "quantized_tensor = linear_q_with_scale_and_zero_point(\n",
    "    test_tensor, scale=scale, zero_point=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "quantized_tensor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Now, put all this in `linear_q_symmetric_per_channel` function defined earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 335
   },
   "outputs": [],
   "source": [
    "def linear_q_symmetric_per_channel(r_tensor, dim, dtype=torch.int8):\n",
    "    \n",
    "    output_dim = r_tensor.shape[dim]\n",
    "    # store the scales\n",
    "    scale = torch.zeros(output_dim)\n",
    "\n",
    "    for index in range(output_dim):\n",
    "        sub_tensor = r_tensor.select(dim, index)\n",
    "        scale[index] = get_q_scale_symmetric(sub_tensor, dtype=dtype)\n",
    "\n",
    "    # reshape the scale\n",
    "    scale_shape = [1] * r_tensor.dim()\n",
    "    scale_shape[dim] = -1\n",
    "    scale = scale.view(scale_shape)\n",
    "    quantized_tensor = linear_q_with_scale_and_zero_point(\n",
    "        r_tensor, scale=scale, zero_point=0, dtype=dtype)\n",
    "   \n",
    "    return quantized_tensor, scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 97
   },
   "outputs": [],
   "source": [
    "test_tensor=torch.tensor(\n",
    "    [[191.6, -13.5, 728.6],\n",
    "     [92.14, 295.5,  -184],\n",
    "     [0,     684.6, 245.5]]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 131
   },
   "outputs": [],
   "source": [
    "### along the rows (dim = 0)\n",
    "quantized_tensor_0, scale_0 = linear_q_symmetric_per_channel(\n",
    "    test_tensor, dim=0)\n",
    "\n",
    "### along the columns (dim = 1)\n",
    "quantized_tensor_1, scale_1 = linear_q_symmetric_per_channel(\n",
    "    test_tensor, dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Plot the quantization error for along the rows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 97
   },
   "outputs": [],
   "source": [
    "dequantized_tensor_0 = linear_dequantization(\n",
    "    quantized_tensor_0, scale_0, 0)\n",
    "\n",
    "plot_quantization_errors(\n",
    "    test_tensor, quantized_tensor_0, dequantized_tensor_0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 46
   },
   "outputs": [],
   "source": [
    "print(f\"\"\"Quantization Error : \\\n",
    "{quantization_error(test_tensor, dequantized_tensor_0)}\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Plot the quantization error for along the columns."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 165
   },
   "outputs": [],
   "source": [
    "dequantized_tensor_1 = linear_dequantization(\n",
    "    quantized_tensor_1, scale_1, 0)\n",
    "\n",
    "plot_quantization_errors(\n",
    "    test_tensor, quantized_tensor_1, dequantized_tensor_1, n_bits=8)\n",
    "\n",
    "print(f\"\"\"Quantization Error : \\\n",
    "{quantization_error(test_tensor, dequantized_tensor_1)}\"\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [
    "NgmjMISuIyzF",
    "3dMVgqcNJCqE",
    "MFx2m7RmzRd5",
    "l6XbkmOzYMrC",
    "8NS1TnQt6E6v"
   ],
   "provenance": [
    {
     "file_id": "12_pQW6LB80u98m72_YKwM4ph7eb5Uf3g",
     "timestamp": 1705360205453
    },
    {
     "file_id": "1U9pm4j_uAD8EO7OrEPvdpwFjQKO3suuH",
     "timestamp": 1700237940920
    }
   ]
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
